
/**
 * https://leetcode.cn/submissions/detail/570259354/
 * 834. 树中距离之和
 * Hard, 黄伟杰 2024.10.05
 * 换根DP
 */

#include <bits/stdc++.h>
using namespace std;

class Solution
{
public:
    void dfs(int x, int fa, int depth, vector<vector<int>> &g, vector<int> &size, vector<int> &ans)
    {                      // 统计每个节点的大小
        ans[0] += depth;   // depth为0到x的距离
        for (int y : g[x]) // 遍历邻居
        {
            if (y != fa)
            {
                dfs(y, x, depth + 1, g, size, ans); // x是y的父节点
                size[x] += size[y];                 // 累加x的儿子y的子树大小
            }
        }
    }

    void reroot(int x, int fa, int n, vector<vector<int>> &g, vector<int> &size, vector<int> &ans)
    { // 计算距离
        for (int y : g[x])
        {
            if (y != fa)
            {
                ans[y] = ans[x] + n - 2 * size[y];
                reroot(y, x, n, g, size, ans);
            }
        }
    }

    vector<int> sumOfDistancesInTree(int n, vector<vector<int>> &edges)
    {
        vector<vector<int>> g(n);
        for (auto &i : edges) // 找到此节点的邻居
        {
            int x = i[0], y = i[1];
            g[x].push_back(y);
            g[y].push_back(x);
        }
        vector<int> ans(n);
        vector<int> size(n, 1); // 每个节点的大小
        dfs(0, -1, 0, g, size, ans);
        reroot(0, -1, n, g, size, ans);

        return ans;
    }
};
